import torch
import torch.nn as nn
import torchvision.models as models
from torchvision.models.resnet import BasicBlock, Bottleneck

class Cifar10ResNet(nn.Module):
    """
    改进的ResNet模型，适用于CIFAR10数据集
    
    特点:
    - 使用更深的ResNet34架构
    - 添加了Dropout层防止过拟合
    - 修改了初始卷积层适应32x32的小图像
    - 添加了批归一化层
    """
    def __init__(self, num_classes=10):
        super(Cifar10ResNet, self).__init__()
        
        # 使用预训练的ResNet34作为基础模型
        self.resnet = models.resnet34(pretrained=False)
        
        # 修改第一层卷积，因为CIFAR10图像较小(32x32)
        self.resnet.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)
        
        # 移除原有的最大池化层，因为图像太小
        self.resnet.maxpool = nn.Identity()
        
        # 添加Dropout层
        self.resnet.fc = nn.Sequential(
            nn.Dropout(0.5),  # 添加Dropout防止过拟合
            nn.Linear(512, num_classes)
        )
        
        # 初始化权重
        self._initialize_weights()
    
    def _initialize_weights(self):
        """
        初始化模型权重
        """
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        return self.resnet(x)